"""
给定一个智能体，环境，然后开始玩游戏
agent
net
"""

from envs.snake.snake import SnakeEnv
from dqn import DQN
import torch
import time
from multiprocessing import Process, Queue, set_start_method
import pygame
from envs.snake import setting
import sys

# 模型的路径
model_path = 'models\dqn_local_search_9000.0.pth'

# 游戏刷新频率, 单位: s
game_fre = 0.1


env = SnakeEnv()
agent = DQN('','','','','','','','', init_type=1, path=model_path)

# pygame 初始化
pygame.init()
screen = pygame.display.set_mode((setting.width * setting.block_size, setting.height * setting.block_size))

pygame.display.set_caption('wjy_snake')
# screen.fill([255, 255, 255])


def convert(s: str):
    """
    输入一个字符，转化成数字形式的动作
    """
    if s == 'w':
        return 0
    elif s == 's':
        return 1
    elif s == 'a':
        return 2
    elif s == 'd':
        return 3
    else:
        return -1


# 这里使用了生产者消费者模型，s是地图
def render(q):
    """
    收到一个状态，就显示出来
    """
    while True:
        # 画矩形
        # lock.wait()
        s = q.get()
        # print(len(s), len(s[0]))
        for i in range(len(s)):
            for j in range(len(s[0])):
                if s[i][j] == setting.wall:
                    color = setting.color_wall
                elif s[i][j] == setting.space:
                    color = setting.color_space
                elif s[i][j] == setting.head:
                    color = setting.color_head
                elif s[i][j] == setting.body:
                    color = setting.color_body
                elif s[i][j] == setting.food:
                    color = setting.color_food
                pygame.draw.rect(screen, color, [i * setting.block_size, j * setting.block_size,
                                                                  setting.block_size, setting.block_size], 0)
        # 循环获取事件，监听事件状态
        for event in pygame.event.get():
            # 判断用户是否点了"X"关闭按钮,并执行if代码段
            if event.type == pygame.QUIT:
                # 卸载所有模块
                pygame.quit()
                # 终止程序，确保退出程序
                sys.exit()
        pygame.display.flip()  # 更新屏幕内容
        # lock.clear()


def run(q, epoch):
    # 所有轮次的奖励
    max_step = 999999
    epoch_reward_list = []
    for _ in range(epoch):
        state = env.reset()
        action = agent.take_action(state)
        done = False
        # 轮次奖励
        epoch_r = 0
        step = 0
        while not done and step < max_step:
            step += 1
            q.put(env.state)
            # lock.set()
            state, reward, done, _ = env.step(action)
            action = agent.take_action(state)
            print(reward)
            epoch_r += reward
            time.sleep(game_fre)
        epoch_reward_list.append(epoch_r)

    for r in epoch_reward_list:
        print(r)




if __name__ == '__main__':
    # set_start_method('fork')

    q=Queue()

    p1=Process(target=run,args=(q, 100))


    c1=Process(target=render,args=(q, ))

    #开始
    p1.start()
    c1.start()
    # print('进程间通信-队列-主进程')